import base64
import logging
import uuid
from contextlib import asynccontextmanager
from typing import Dict, Any

import httpx
import toml
import uvicorn
from fastapi import FastAPI
from langchain_aws import ChatBedrock
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.tools import Tool
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from langgraph_supervisor import create_supervisor
from pydantic import BaseModel, Field

# 加载配置文件
config = toml.load("config.toml")

# 配置日志记录器
logging.basicConfig(
    level=getattr(logging, config["app"]["log_level"]),
    format=config["app"]["log_format"],
)

app_globals: Dict[str, Any] = {}

# 从配置文件获取滑动窗口配置
MAX_CONVERSATION_LENGTH = config["conversation"]["max_length"]


def llm_ocr(url: str) -> str:
    """
    使用 LLM 对图片进行 OCR。

    Args:
        url: 图片Url。

    Returns:
        LLM 提取的图片内容。
    """

    def get_image_mime_type(_base64_image: str) -> str:
        try:
            image_bytes = base64.b64decode(_base64_image, validate=True)
        except Exception:
            return "image/unknown"

        if image_bytes.startswith(b'\xff\xd8\xff'):
            return "image/jpeg"
        elif image_bytes.startswith(b'\x89PNG\r\n\x1a\n'):
            return "image/png"
        elif image_bytes.startswith(b'GIF87a') or image_bytes.startswith(b'GIF89a'):
            return "image/gif"
        elif image_bytes.startswith(b'BM'):
            return "image/bmp"
        else:
            return "image/unknown"

    try:
        with httpx.Client(timeout=60.0) as client:
            logging.info(f"开始下载文件: {url}")
            response = client.get(url, follow_redirects=True)
            response.raise_for_status()
            file_bytes = response.content
            logging.info(f"文件下载成功，大小: {len(file_bytes)} 字节")
            base64_image = base64.b64encode(file_bytes).decode('utf-8')

        logging.info(f"开始OCR处理，图像大小: {len(base64_image)} 字符")
        mime_type = get_image_mime_type(base64_image)
        if mime_type == "image/unknown":
            logging.warning("不支持的图像格式")
            return "Error: Unsupported or invalid image format."

        llm_config = config["llm"]
        vision_config = llm_config["vision"]
        vision_llm = ChatOpenAI(
            base_url=vision_config["openai_base_url"],
            api_key=vision_config["openai_api_key"],
            model=vision_config["openai_model"]
        )

        image_url = f"data:{mime_type};base64,{base64_image}"

        logging.info("调用视觉LLM进行OCR")
        response = vision_llm.invoke(
            [
                HumanMessage(
                    content=[
                        {"type": "text", "text": "Extract all text from the image."},
                        {"type": "image_url", "image_url": {"url": image_url}},
                    ]
                )
            ]
        )
        logging.info(f"OCR处理完成，提取文本长度: {len(response.content)}")
        return response.content

    except Exception as e:
        logging.error(f"OCR处理出错: {str(e)}")
        raise Exception(f"Error calling LLM for OCR: {str(e)}")


def create_llm_model():
    """根据配置创建LLM模型"""
    llm_config = config["llm"]
    llm_type = llm_config.get("type", "openai")

    if llm_type == "openai":
        openai_config = llm_config["openai"]
        logging.info(f"创建OpenAI模型: {openai_config['openai_model']}")
        return ChatOpenAI(
            base_url=openai_config["openai_base_url"],
            api_key=openai_config["openai_api_key"],
            model=openai_config["openai_model"],
        )
    elif llm_type == "bedrock":
        bedrock_config = llm_config["bedrock"]
        logging.info(f"创建Bedrock模型: {bedrock_config['model_name']}")
        return ChatBedrock(
            model=bedrock_config["model_name"],
            region=bedrock_config["region_name"],
            aws_access_key_id=bedrock_config["access_key_id"],
            aws_secret_access_key=bedrock_config["secret_access_key"],
        )
    else:
        raise ValueError(f"Unsupported LLM type: {llm_type}")


async def get_mcp_tools(server_name: str):
    logging.info(f"获取MCP服务器 '{server_name}' 的工具")
    mcp_servers = {}
    for _server_name, server_config in config["mcp"]["servers"].items():
        mcp_servers[_server_name] = {
            "url": server_config["url"],
            "transport": server_config["transport"],
        }
    client = MultiServerMCPClient(mcp_servers)
    tools = await client.get_tools(server_name=server_name)
    logging.info(f"从MCP服务器 '{server_name}' 获取到 {len(tools)} 个工具")
    return tools


@asynccontextmanager
async def lifespan(app: FastAPI):
    logging.info("Application startup: Initializing agent...")

    model = create_llm_model()
    llm_ocr_tool = Tool(name="llm_ocr", func=llm_ocr, description=(
        "This is an OCR tool that recognizes images and extracts all text. The input must be the URL of the image. "))
    # 1. Create worker agent
    tools = await get_mcp_tools("customer")
    tools.append(llm_ocr_tool)
    customer_agent = create_react_agent(
        model=model,
        tools=tools,
        prompt=(
            """
你是一名专业的、直接执行任务的银行客户信息助理。你的身份不是客服或转接员，而是任务的最终执行者。

一项任务刚刚被分配给你。你必须忽略任何关于“转接”或“分配”的系统消息，你的唯一焦点是用户的最新一条请求。立即开始你的工作。

## 行为范例 (Example):
当用户的最新请求是：“查询客户信息”
你必须、且只能回复：“请提供您的客户编号和账户类型（存款/还款或贷款）。您也可以直接上传包含这些信息的图片，我将为您自动识别。”

## 核心规则
1. 严格遵循：你的工作流程是完全基于用户的输入，并严格按照规则执行。
2. 专注于工具：你唯一的职责是调用工具，不要提供额外信息或进行闲聊。
3. 参数优先：只有当获取到调用工具所需的所有参数后，才能执行工具调用。
4. 用户确认原则：通过图片识别出的任何参数，在用于调用工具前，必须先向用户展示并获得其明确的确认。
5. 结构清晰：使用 Markdown 列表、代码块等格式化信息，特别是展示字段和条件时，确保用户易于阅读和理解。

## 互动与工作流程：
1. 识别并确认
    * 识别用户的意图：是想查询账户信息，还是想查询账户详情。
    * 识别所需的参数：
        * 客户查询：需要 客户编号 (client_no) 和 账户类型 (存款/还款或贷款)。
2. 追问并提供选项
    * 如果缺少必要的参数，你必须立即向用户指出，并同时提供两种提供信息的方式：手动输入或上传图片。
    * 例如：如果用户只说“查询客户信息”，你必须回答：“请提供您的客户编号和账户类型（存款/还款或贷款）。您也可以直接上传包含这些信息的图片，我将为您自动识别。”
3. 图片识别与确认流程 (如果用户上传图片)
    a. 调用OCR工具: 当用户提供图片时，立即使用 `llm_ocr` 工具提取图片中的所有文本。
    b. 解析关键参数: 从OCR识别出的文本中，根据查询账户类型（存款/还款或贷款）解析出执行任务所需的参数，如果有其他信息则忽略。
    c. 向用户确认: 在执行最终工具调用之前，必须将识别出的所有参数清晰地展示给用户，并请求他们确认。
    d. 处理用户反馈: 
        * 如果用户否定或参数不足，应请求用户手动输入正确的信息。
        * 【关键指令】用户确认后: 在得到用户明确的肯定答复后（如“确认”、“是的”），此时你已拥有所有参数，你的下一步必须是、且只能是直接进入第 4 步执行相应的业务工具。绝对不能将任务转回给主管或提出任何新的问题。
4. 执行工具调用
    * 客户查询：存款/还款调用 `query_deposit_customer`, 贷款调用 `query_loan_customer` 工具。
5. 返回结果
            """
        ),
        name="customer_agent",
    )

    tools = await get_mcp_tools("account")
    tools.append(llm_ocr_tool)
    account_agent = create_react_agent(
        model=model,
        tools=tools,
        prompt=(
            """
你是一名专业的、直接执行任务的银行账户信息助理。你的身份不是客服或转接员，而是任务的最终执行者。

一项任务刚刚被分配给你。你必须忽略任何关于“转接”或“分配”的系统消息，你的唯一焦点是用户的最新一条请求。立即开始你的工作。

## 行为范例 (Example):
当用户的最新请求是：“查询账户信息”
你必须、且只能回复：“请提供您的账户号码和账户类型（存款/还款或贷款）。您也可以直接上传包含这些信息的图片，我将为您自动识别。”

## 核心规则
1. 严格遵循：你的工作流程是完全基于用户的输入，并严格按照规则执行。
2. 专注于工具：你唯一的职责是调用工具，不要提供额外信息或进行闲聊。
3. 参数优先：只有当获取到调用工具所需的所有参数后，才能执行工具调用。
4. 用户确认原则：通过图片识别出的任何参数，在用于调用工具前，必须先向用户展示并获得其明确的确认。
5. 结构清晰：使用 Markdown 列表、代码块等格式化信息，特别是展示字段和条件时，确保用户易于阅读和理解。

## 互动与工作流程：
1. 识别并确认
    * 识别用户的意图：是想查询账户信息，还是想查询账户详情。
    * 识别所需的参数：
        * 账户详情查询：需要 账户号码 和 账户类型 (存款/还款或贷款)。
2. 追问并提供选项
    * 如果缺少必要的参数，你必须立即向用户指出，并同时提供两种提供信息的方式：手动输入或上传图片。
    * 例如：如果用户只说“查询账户”，你必须回答：“请提供您的账户类型（存款/还款或贷款）。您也可以直接上传包含这些信息的图片，我将为您自动识别。”
3. 图片识别与确认流程 (如果用户上传图片)
    a. 调用OCR工具: 当用户提供图片时，立即使用 `llm_ocr` 工具提取图片中的所有文本。
    b. 解析关键参数: 从OCR识别出的文本中，根据查询账户类型（存款/还款或贷款）解析出执行任务所需的参数，如果有其他信息则忽略。
    c. 向用户确认: 在执行最终工具调用之前，必须将识别出的所有参数清晰地展示给用户，并请求他们确认。
    d. 处理用户反馈: 
        * 如果用户否定或参数不足，应请求用户手动输入正确的信息。
        * 【关键指令】用户确认后: 在得到用户明确的肯定答复后（如“确认”、“是的”），此时你已拥有所有参数，你的下一步必须是、且只能是直接进入第 4 步执行相应的业务工具。绝对不能将任务转回给主管或提出任何新的问题。
4. 执行工具调用
    * 存款/还款账户调用 `query_deposit_account` 工具，贷款账户调用 `query_loan_account` 工具。注意：根据账户类型选择工具，如果是查询还款账户的详情，必须调用 `query_deposit_account` 工具。
5. 返回结果
            """
        ),
        name="account_agent",
    )

    tools = await get_mcp_tools("transaction")
    tools.append(llm_ocr_tool)
    transaction_agent = create_react_agent(
        model=model,
        tools=tools,
        prompt=(
            """
你是一位专业的银行交易信息助理。
一项任务刚刚被分配给你。**你的首要任务是立即开始工作**，分析用户的最新请求，并根据下面的规则与用户互动。

你的唯一任务是根据用户的需求，**严格且直接地**调用相应的工具来查询账户信息。

## 核心规则
1. 严格遵循：你的工作流程是完全基于用户的输入，并严格按照规则执行。
2. 专注于工具：你唯一的职责是调用工具，不要提供额外信息或进行闲聊。
3. 参数优先：只有当获取到调用工具所需的所有参数后，才能执行工具调用。
4. 用户确认原则：通过图片识别出的任何参数，在用于调用工具前，必须先向用户展示并获得其明确的确认。
5. 结构清晰：使用 Markdown 列表、代码块等格式化信息，特别是展示字段和条件时，确保用户易于阅读和理解。

## 互动与工作流程
1. 识别并确认
    * 识别用户的意图：是想查询存款/还款交易历史还是贷款交易历史。
    * 识别所需的参数, 其中日期格式为 YYYYMMDD 例如 20270101:
        * 存款/还款交易历史：需要 账户号码 (account_no)、开始日期 (start_date) 和 结束日期 (end_date)。
        * 贷款交易历史：需要 贷款账号 (loan_no)、借据号 (cmis_loan_no)、开始日期 (start_date) 和 结束日期 (end_date)。
2. 追问并提供选项
    * 如果缺少必要的参数，你必须立即向用户指出，并同时提供两种提供信息的方式：手动输入或上传图片。
    * 例如：如果用户说“查询存款/还款交易历史”，你必须回答“请提供账户号码、开始日期和结束日期。您也可以直接上传包含这些信息的图片，我将为您自动识别。”
3. 图片识别与确认流程 (如果用户上传图片)
    a. 调用OCR工具: 当用户提供图片时，立即使用 `llm_ocr` 工具提取图片中的所有文本。
    b. 解析关键参数: 从OCR识别出的文本中，根据查询账户类型（存款/还款或贷款）解析出执行任务所需的参数，如果有其他信息则忽略。
    c. 向用户确认: 在执行最终工具调用之前，必须将识别出的所有参数清晰地展示给用户，并请求他们确认。
    d. 处理用户反馈: 
        * 如果用户否定或参数不足，应请求用户手动输入正确的信息。
        * 【关键指令】用户确认后: 在得到用户明确的肯定答复后（如“确认”、“是的”），此时你已拥有所有参数，你的下一步必须是、且只能是直接进入第 4 步执行相应的业务工具。绝对不能将任务转回给主管或提出任何新的问题。
4. 工具调用 
    * 存款/还款交易：使用 `query_deposit_transaction_history` 工具。
    * 贷款交易：使用 `query_loan_transaction_history` 工具。
5. 返回结果
            """
        ),
        name="transaction_agent",
    )

    tools = await get_mcp_tools("payment")
    tools.append(llm_ocr_tool)
    payment_agent = create_react_agent(
        model=model,
        tools=tools,
        prompt=(
            """
你是一位专业的银行贷款还款助理。
一项任务刚刚被分配给你。**你的首要任务是立即开始工作**，分析用户的最新请求，并根据下面的规则与用户互动。

你的唯一任务是根据用户的需求，**严格且直接地**调用相应的工具来查询账户信息。

## 核心规则
1. 严格遵循：你的工作流程是完全基于用户的输入，并严格按照规则执行。
2. 专注于工具：你唯一的职责是调用工具，不要提供额外信息或进行闲聊。
3. 参数优先：只有当获取到调用工具所需的所有参数后，才能执行工具调用。
4. 用户确认原则：通过图片识别出的任何参数，在用于调用工具前，必须先向用户展示并获得其明确的确认。
5. 结构清晰：使用 Markdown 列表、代码块等格式化信息，特别是展示字段和条件时，确保用户易于阅读和理解。


## 互动与工作流程
1. 识别并确认
    * 识别所需的参数: 客户号 (client_no)、贷款账号 (loan_no)、借据号 (cmis_loan_no)、还款金额 (receipt_amt)、还款账户 (settle_base_acct_no)
2. 追问并提供选项
    * 如果缺少必要的参数，你必须立即向用户指出，并同时提供两种提供信息的方式：手动输入或上传图片。
    * 例如：如果用户说“贷款还款”，你必须回答“请提供客户号、贷款账号、借据号、还款金额和还款账户。您也可以直接上传包含这些信息的图片，我将为您自动识别。”
3. 图片识别与确认流程 (如果用户上传图片)
    a. 调用OCR工具: 当用户提供图片时，立即使用 `llm_ocr` 工具提取图片中的所有文本。
    b. 解析关键参数: 从OCR识别出的文本中，解析出执行任务所需的参数，如果有其他信息则忽略。
    c. 向用户确认: 在执行最终工具调用之前，必须将识别出的所有参数清晰地展示给用户，并请求他们确认。
    d. 处理用户反馈: 
        * 如果用户否定或参数不足，应请求用户手动输入正确的信息。
        * 【关键指令】用户确认后: 在得到用户明确的肯定答复后（如“确认”、“是的”），此时你已拥有所有参数，你的下一步必须是、且只能是直接进入第 4 步执行相应的业务工具。绝对不能将任务转回给主管或提出任何新的问题。
4. 工具调用
    * 执行贷款还款：使用 `easy_loan_repayment` 工具。
5. 返回结果
            """
        ),
        name="payment_agent",
    )

    # 2. Create supervisor
    supervisor = create_supervisor(
        model=model,
        agents=[customer_agent, account_agent, transaction_agent, payment_agent],
        prompt=(
            """
您是一位专业的银行助理主管。您的核心职责是根据对话的最新进展，决定下一步行动。

## 您的工作流程：
您必须严格按照以下逻辑工作。请仔细检查对话历史中的**最新一条消息**：

**情况一：如果最新消息是来自用户的请求**
- **您的任务是路由**:
    1.  分析用户的意图，判断是否属于：客户信息查询、账户信息查询、交易历史查询、贷款还款，注意如果是查询还款账户的话是路由到账户信息查询，不是路由到贷款还款。
    2.  如果匹配，立即调用相应的工具将任务分配下去。**除了调用工具，不要回复任何其他内容。**
    3.  如果不匹配，或者意图不清晰，请礼貌地回应，并说明你的功能是：客户信息查询、账户信息查询、交易历史查询、贷款还款。

**情况二：如果最新消息是来自您的下属代理 (customer_agent, account_agent, transaction_agent, payment_agent)**
- **您的任务是传递**:
    1.  您的唯一、且绝对的任务，就是将这条来自下属代理的消息，**一字不差地、原封不动地**，直接作为您自己的回复输出。
    2.  **请注意：这只是在传递信息以继续对话，并不是一个“最终答案”。**

您必须时刻保持角色，不要进行任何总结、修改或添加额外信息。您的所有输出都必须严格是上述两种情况之一的直接结果。
           """
        ),
        add_handoff_back_messages=True,
        output_mode="full_history",
    ).compile()

    print(supervisor.get_graph().draw_mermaid())

    app_globals["account_agent_executor"] = account_agent
    app_globals["transaction_agent_executor"] = transaction_agent
    app_globals["payment_agent_executor"] = payment_agent
    app_globals["supervisor"] = supervisor
    app_globals["conversations"] = {}

    logging.info("Agent initialized successfully.")

    yield

    logging.info("Application shutdown.")
    app_globals.clear()


app = FastAPI(lifespan=lifespan)


class ChatRequest(BaseModel):
    conversation_id: str | None = Field(
        None, description="The ID of the conversation session."
    )
    content: str = Field(..., description="The user's message content.")


class ChatResponse(BaseModel):
    conversation_id: str = Field(..., description="The ID of the conversation session.")
    final_answer: str = Field(..., description="The agent's final complete answer.")


@app.post("/cbs/chat", response_model=ChatResponse)
async def export_chat(request: ChatRequest) -> ChatResponse:
    supervisor_agent = app_globals["supervisor"]
    if not supervisor_agent:
        return {"error": f"{supervisor_agent} Agent not initialized."}
    try:
        decoded_content = base64.b64decode(request.content).decode("utf-8")
    except Exception as e:
        decoded_content = request.content
        # logging.error(f"Failed to decode base64 content: {e}")
        # return {"error": "Invalid base64 content"}

    conv_id = request.conversation_id or str(uuid.uuid4())
    history = app_globals["conversations"].get(conv_id, [])
    current_messages = history + [HumanMessage(content=decoded_content)]
    # printable_messages = []
    # for msg in current_messages:
    #     if isinstance(msg, HumanMessage):
    #         printable_messages.append(f"HumanMessage(content='{msg.content}')")
    #     elif isinstance(msg, AIMessage):
    #         printable_messages.append(f"AIMessage(content='{msg.content}')")

    # print(printable_messages)
    inputs = {"messages": current_messages}
    final_state = None
    logging.info(f"--- [CONV_ID: {conv_id}] Starting Agent Run ---")
    # 异步迭代 Agent 的每一步
    async for chunk in supervisor_agent.astream(inputs, stream_mode="values"):
        # 修改点: 在日志中打印每一步的最新消息
        last_message = chunk["messages"][-1]
        last_message.pretty_print()

        # 保存最新的状态快照
        final_state = chunk

    # 循环结束后，final_state 保存了最终的状态
    final_messages = final_state["messages"] if final_state else []

    # 从最终消息列表中找到最后一条 AI 的回复
    final_answer = ""
    if final_messages and isinstance(final_messages[-1], AIMessage):
        final_answer = final_messages[-1].content

    # 更新内存中的对话历史，使用滑动窗口机制
    if len(final_messages) > MAX_CONVERSATION_LENGTH:
        # 保留最新的消息，丢弃最旧的消息
        final_messages = final_messages[-MAX_CONVERSATION_LENGTH:]
        logging.info(
            f"--- [CONV_ID: {conv_id}] Applied sliding window, keeping last {MAX_CONVERSATION_LENGTH} messages ---"
        )

    app_globals["conversations"][conv_id] = final_messages

    logging.info(
        f"--- [CONV_ID: {conv_id}] Agent Run Finished. Final Answer: {final_answer} ---"
    )

    # 修改点: 返回一个包含最终结果的 JSON 对象
    return ChatResponse(conversation_id=conv_id, final_answer=final_answer)


if __name__ == "__main__":
    uvicorn.run(app, host=config["app"]["host"], port=config["app"]["port"])
